import jieba
import sys
import pickle

if __name__ == '__main__':
    test_file = sys.argv[1]
    out_file = sys.argv[2]

    # count_vector_file = sys.argv[1]
    # tfidf_file = sys.argv[2]
    # lg_file = sys.argv[3]

    count_vector_file = '/Users/hardy/data/models/cv20_yz.pkl'
    tfidf_file = '/Users/hardy/data/models/tfidf20_yz.pkl'
    lg_file = '/Users/hardy/data/models/lg20_yz.pkl'

    CV = pickle.load(open(count_vector_file, 'rb'))
    TFIDF = pickle.load(open(tfidf_file, 'rb'))
    LG = pickle.load(open(lg_file, 'rb'))

    with open(out_file, 'w+') as fd:
        for l in open(test_file):
            d = l.strip().split('\t')
            if len(d) != 2:
                continue
            label, text = d
            seg = jieba.lcut(text)

            c = LG.predict(TFIDF.transform(CV.transform([' '.join(seg)])))[0]
            prob = dict(zip(LG.classes_, LG.predict_proba(TFIDF.transform(CV.transform([' '.join(seg)])))[0]))[c]
            if label.strip() == c.strip():
                t = '\t'.join([label, str(prob), text]) + '\n'
                fd.write(t)
